import numpy as np
import torch
from config import *
from pathfinding.core.grid import Grid
from pathfinding.finder.a_star import AStarFinder

from graph_utils import create_distance_maps, create_direction_maps


def create_advanced_state_representation(grid_state):
    """创建高级状态表示，包含路径规划和行为特征"""
    # 基本表示
    basic_features = np.zeros((ENV_SIZE, ENV_SIZE, CELL_FEATURES))
    for i in range(ENV_SIZE):
        for j in range(ENV_SIZE):
            basic_features[i, j, grid_state[i, j]] = 1
            
    # 计算常规距离/方向图
    distance_maps = create_distance_maps(grid_state)
    direction_maps = create_direction_maps(grid_state)
    
    # 创建最短路径地图
    path_map = create_shortest_path_map(grid_state)
    
    # 创建动作影响热图
    action_heatmap = create_action_influence_map(grid_state)
    
    # 创建可达性地图
    reachability_map = create_reachability_map(grid_state)
    
    # 创建目标导向特征
    goal_features = create_goal_directed_features(grid_state)
    
    # 合并所有特征
    all_features = [
        basic_features,
        np.stack(distance_maps, axis=-1),
        np.stack(direction_maps, axis=-1),
        path_map[..., np.newaxis],
        action_heatmap,
        reachability_map[..., np.newaxis],
        goal_features
    ]
    
    # 连接所有特征
    combined_features = np.concatenate([feat.reshape(ENV_SIZE, ENV_SIZE, -1) for feat in all_features], axis=-1)
    
    return combined_features

def create_shortest_path_map(grid_state):
    """创建从智能体到目标的最短路径地图"""
    # 找到智能体和目标位置
    agent_pos = np.where(grid_state == AGENT)
    goal_pos = np.where(grid_state == END)
    
    # 如果智能体或目标不存在，返回空地图
    if len(agent_pos[0]) == 0 or len(goal_pos[0]) == 0:
        return np.zeros((ENV_SIZE, ENV_SIZE))
    
    agent_i, agent_j = agent_pos[0][0], agent_pos[1][0]
    goal_i, goal_j = goal_pos[0][0], goal_pos[1][0]
    
    # 创建寻路网格
    matrix = np.ones((ENV_SIZE, ENV_SIZE))
    matrix[grid_state == OBSTACLE] = 0  # 障碍物不可通行
    
    grid = Grid(matrix=matrix)
    start = grid.node(agent_j, agent_i)  # 注意：Grid使用(x,y)而不是(i,j)
    end = grid.node(goal_j, goal_i)
    
    # 寻找路径
    finder = AStarFinder()
    path, _ = finder.find_path(start, end, grid)
    
    # 创建路径地图
    path_map = np.zeros((ENV_SIZE, ENV_SIZE))
    for x, y in path:
        path_map[y, x] = 1  # 将坐标转回(i,j)
    
    return path_map

def create_action_influence_map(grid_state):
    """创建动作影响热图，表示每个动作的可行性和价值"""
    # 为每个动作创建一个热图层
    heatmap = np.zeros((ENV_SIZE, ENV_SIZE, 5))  # 5个动作：上下左右停留
    
    # 找到智能体位置
    agent_pos = np.where(grid_state == AGENT)
    if len(agent_pos[0]) == 0:
        return heatmap
    
    agent_i, agent_j = agent_pos[0][0], agent_pos[1][0]
    
    # 为每个位置计算每个动作的影响
    for i in range(ENV_SIZE):
        for j in range(ENV_SIZE):
            # 上
            if i > 0 and grid_state[i-1, j] != OBSTACLE:
                heatmap[i, j, 0] = 1
            
            # 下
            if i < ENV_SIZE-1 and grid_state[i+1, j] != OBSTACLE:
                heatmap[i, j, 1] = 1
                
            # 左
            if j > 0 and grid_state[i, j-1] != OBSTACLE:
                heatmap[i, j, 2] = 1
                
            # 右
            if j < ENV_SIZE-1 and grid_state[i, j+1] != OBSTACLE:
                heatmap[i, j, 3] = 1
                
            # 停留
            if grid_state[i, j] != OBSTACLE:
                heatmap[i, j, 4] = 1
    
    return heatmap

def create_reachability_map(grid_state, max_steps=5):
    """创建可达性地图，表示从智能体出发在n步内可以到达的位置"""
    # 找到智能体位置
    agent_pos = np.where(grid_state == AGENT)
    if len(agent_pos[0]) == 0:
        return np.zeros((ENV_SIZE, ENV_SIZE))
    
    agent_i, agent_j = agent_pos[0][0], agent_pos[1][0]
    
    # 初始化可达性地图
    reachability = np.zeros((ENV_SIZE, ENV_SIZE))
    reachability[agent_i, agent_j] = 1
    
    # 障碍物掩码
    obstacle_mask = (grid_state == OBSTACLE)
    
    # BFS扩散
    queue = [(agent_i, agent_j, 0)]  # (i, j, steps)
    visited = set([(agent_i, agent_j)])
    
    while queue:
        i, j, steps = queue.pop(0)
        
        if steps >= max_steps:
            continue
            
        # 检查四个方向
        for di, dj in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
            ni, nj = i + di, j + dj
            
            # 检查边界和障碍物
            if (0 <= ni < ENV_SIZE and 0 <= nj < ENV_SIZE and 
                not obstacle_mask[ni, nj] and (ni, nj) not in visited):
                
                # 标记为可达
                reachability[ni, nj] = 1 - (steps + 1) / (max_steps + 1)  # 距离越远，值越小
                
                # 添加到队列和已访问集合
                queue.append((ni, nj, steps + 1))
                visited.add((ni, nj))
    
    return reachability

def create_goal_directed_features(grid_state):
    """创建目标导向特征，表示朝向目标移动的方向和距离"""
    # 找到智能体和目标位置
    agent_pos = np.where(grid_state == AGENT)
    goal_pos = np.where(grid_state == END)
    
    # 如果智能体或目标不存在，返回空特征
    if len(agent_pos[0]) == 0 or len(goal_pos[0]) == 0:
        return np.zeros((ENV_SIZE, ENV_SIZE, 4))
    
    agent_i, agent_j = agent_pos[0][0], agent_pos[1][0]
    goal_i, goal_j = goal_pos[0][0], goal_pos[1][0]
    
    # 计算目标方向向量
    di = goal_i - agent_i
    dj = goal_j - agent_j
    
    # 计算曼哈顿距离
    manhattan_dist = abs(di) + abs(dj)
    
    # 归一化方向
    direction = np.zeros(4)  # 上下左右的方向分量
    
    if manhattan_dist > 0:
        if di < 0:
            direction[0] = abs(di) / manhattan_dist  # 上
        elif di > 0:
            direction[1] = abs(di) / manhattan_dist  # 下
            
        if dj < 0:
            direction[2] = abs(dj) / manhattan_dist  # 左
        elif dj > 0:
            direction[3] = abs(dj) / manhattan_dist  # 右
    
    # 为每个位置创建目标特征
    goal_features = np.zeros((ENV_SIZE, ENV_SIZE, 4))
    for i in range(ENV_SIZE):
        for j in range(ENV_SIZE):
            # 对于智能体位置，使用方向向量
            if i == agent_i and j == agent_j:
                goal_features[i, j] = direction
                
            # 对于其他位置，根据相对位置计算
            else:
                # 计算相对于智能体的方向
                rel_di = i - agent_i
                rel_dj = j - agent_j
                rel_dist = abs(rel_di) + abs(rel_dj)
                
                if rel_dist > 0:
                    # 计算该位置是否在通向目标的方向上
                    alignment = (rel_di * di + rel_dj * dj) / (rel_dist * manhattan_dist) if manhattan_dist > 0 else 0
                    
                    # 标记在目标方向上的位置
                    if alignment > 0:
                        if rel_di < 0 and di < 0:
                            goal_features[i, j, 0] = alignment  # 上
                        elif rel_di > 0 and di > 0:
                            goal_features[i, j, 1] = alignment  # 下
                            
                        if rel_dj < 0 and dj < 0:
                            goal_features[i, j, 2] = alignment  # 左
                        elif rel_dj > 0 and dj > 0:
                            goal_features[i, j, 3] = alignment  # 右
    
    return goal_features